package com.example.demo.service.impl;

import com.example.demo.config.LocalRateLimitConfig;
import com.example.demo.service.ILocalRateLimiter;
import com.example.demo.util.WebFluxUtil;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * 本地内存实现对 ip 限流
 * 简单的使用缓存对 ip 进行限流
 *
 * @author zxb
 **/
@Service(value = LocalRateLimitConfig.Constants.IP_BEAN_NAME)
@Slf4j
public class LocalIpRateLimitSimpleImpl implements ILocalRateLimiter {

    /**
     * ConcurrentHashMap 的替代方式，实现自动释放内存
     */
    private LoadingCache<String, RateLimiter> rateLimiterLoadingCache;

    @Autowired
    private LocalRateLimitConfig.LocalRateLimitProperties localRateLimitProperties;

    @PostConstruct
    public void init() {
        Map<String, Object> extend = localRateLimitProperties.getExtend();
        if (Objects.isNull(extend) || Objects.isNull(extend.get(LocalRateLimitConfig.Constants.IP_EXPIRE_AFTER_ACCESS_FIELD_NAME))) {
            throw new IllegalArgumentException("未配置扩展参数-ip 缓存回收时间");
        }
        long expireAfterAccess = Long.parseLong(extend.get(LocalRateLimitConfig.Constants.IP_EXPIRE_AFTER_ACCESS_FIELD_NAME).toString());

        // 创建每秒 n 个令牌桶，如果 30 分钟内没有进行读写则回收
        rateLimiterLoadingCache = CacheBuilder.newBuilder()
                .expireAfterAccess(expireAfterAccess, TimeUnit.SECONDS)
                .build(new CacheLoader<>() {
                    @Override
                    public RateLimiter load(String key) {
                        // 平滑突发限流
                        return RateLimiter.create(localRateLimitProperties.getPermitsPerSecond());
                    }
                });
    }

    @Override
    public RateLimiterResponse access(Object key) {
        log.info("请求的ip {}", getLimitKey(key));
        String ip = getLimitKey(key);
        RateLimiter rateLimiter = rateLimiterLoadingCache.getUnchecked(ip);
        return new RateLimiterResponse(rateLimiter.tryAcquire());
    }

    @Override
    public String getLimitKey(Object key) {
        if (key instanceof ServerHttpRequest) {
            return WebFluxUtil.getIpAddress((ServerHttpRequest) key);
        }
        throw new UnsupportedOperationException();
    }
}
